package prefuse.data.search;

import java.util.Iterator;
import java.util.LinkedList;
import java.util.NoSuchElementException;

import prefuse.data.Tuple;

/**
 * A trie data structure for fast-lookup of words based on their
 * prefixes. The name "Trie" is a play on the words "tree" and 
 * "retrieval". This class builds a tree structure representing a set of
 * words by their prefixes. It is useful for performing prefix-based
 * searches over large amounts of text in an efficient manner.
 *
 * @version 1.0
 * @author <a href="http://jheer.org">jeffrey heer</a>
 * @see PrefixSearchTupleSet
 */
public class Trie {

    /**
     * Base class for nodes in the trie structure.
     */
    public class TrieNode {
        boolean isLeaf;
        int leafCount = 0;
    }

    /**
     * A TrieNode implementation representing a branch in the tree. The
     * class maintains a list of characters (the next character in the
     * prefix) and associated children TrieNodes for each.
     */
    public class TrieBranch extends TrieNode {
        char[] chars = new char[] { 0 };
        TrieNode[] children = new TrieNode[1];
    }

    /**
     * A TrieNode implementation representing a leaf in the tree. The class
     * stores the word and tuple for the leaf, as well as a reference to the
     * successor leaf node in the trie.
     */
    public class TrieLeaf extends TrieNode {
        public TrieLeaf(String word, Tuple t) {
            this.word = word;
            tuple = t;
            next = null;
            leafCount = 1;
        }

        String word;
        Tuple tuple;
        TrieLeaf next;
    }

    /**
     * An iterator for traversing a subtree of the Trie.
     */
    public class TrieIterator implements Iterator {
        private LinkedList queue;

        public TrieIterator(TrieNode node) {
            queue = new LinkedList();
            queue.add(node);
        }

        public boolean hasNext() {
            return !queue.isEmpty();
        }

        public Object next() {
            if (queue.isEmpty())
                throw new NoSuchElementException();

            TrieNode n = (TrieNode) queue.removeFirst();
            Object o;
            if (n instanceof TrieLeaf) {
                TrieLeaf l = (TrieLeaf) n;
                o = l.tuple;
                if (l.next != null)
                    queue.addFirst(l.next);
                return o;
            } else {
                TrieBranch b = (TrieBranch) n;
                for (int i = b.chars.length - 1; i > 0; i--) {
                    queue.addFirst(b.children[i]);
                }
                if (b.children[0] != null)
                    queue.addFirst(b.children[0]);
                return next();
            }
        }

        public void remove() {
            throw new UnsupportedOperationException();
        }
    } // end of inner clas TrieIterator

    private TrieBranch root = new TrieBranch();
    private boolean caseSensitive = false;

    /**
     * Create a new Trie with the specified case-sensitivity.
     * @param caseSensitive true if the index should be case sensitive for
     * indexed words, false otherwise.
     */
    public Trie(boolean caseSensitive) {
        this.caseSensitive = caseSensitive;
    }

    /**
     * Indicates if this Trie's index takes the case of letters
     * into account.
     * @return true if the index is case-sensitive, false otherwise
     */
    public boolean isCaseSensitive() {
        return caseSensitive;
    }

    /**
     * Add a new word to the trie, associated with the given Tuple.
     * @param word the word to add to the Trie
     * @param t the Tuple associated with the word
     */
    public void addString(String word, Tuple t) {
        TrieLeaf leaf = new TrieLeaf(word, t);
        addLeaf(root, leaf, 0);
    }

    /**
     * Remove a word/Tuple pair from the trie.
     * @param word the word to remove
     * @param t the associate Tuple to remove
     */
    public void removeString(String word, Tuple t) {
        removeLeaf(root, word, t, 0);
    }

    private final int getIndex(char[] chars, char c) {
        for (int i = 0; i < chars.length; i++)
            if (chars[i] == c)
                return i;
        return -1;
    }

    private final char getChar(String s, int i) {
        char c = (i < 0 || i >= s.length() ? 0 : s.charAt(i));
        return (caseSensitive ? c : Character.toLowerCase(c));
    }

    private final TrieNode equalityCheck(String word, TrieLeaf l) {
        if (caseSensitive) {
            return l.word.startsWith(word) ? l : null;
        } else {
            // do our own looping to avoid string allocation for case change
            int len = word.length();
            if (len > l.word.length())
                return null;
            for (int i = 0; i < len; ++i) {
                char c1 = Character.toLowerCase(word.charAt(i));
                char c2 = Character.toLowerCase(l.word.charAt(i));
                if (c1 != c2)
                    return null;
            }
            return l;
        }
    }

    private boolean removeLeaf(TrieBranch b, String word, Tuple t, int depth) {
        char c = getChar(word, depth);
        int i = getIndex(b.chars, c);

        if (i == -1) {
            // couldn't find leaf
            return false;
        } else {
            TrieNode n = b.children[i];
            if (n instanceof TrieBranch) {
                TrieBranch tb = (TrieBranch) n;
                boolean rem = removeLeaf(tb, word, t, depth + 1);
                if (rem) {
                    b.leafCount--;
                    if (tb.leafCount == 1)
                        b.children[i] = tb.children[tb.children[0] != null ? 0 : 1];
                }
                return rem;
            } else {
                TrieLeaf nl = (TrieLeaf) n;
                if (nl.tuple == t) {
                    b.children[i] = nl.next;
                    if (nl.next == null)
                        repairBranch(b, i);
                    b.leafCount--;
                    return true;
                } else {
                    TrieLeaf nnl = nl.next;
                    while (nnl != null && nnl.tuple != t) {
                        nl = nnl;
                        nnl = nnl.next;
                    }
                    if (nnl == null)
                        return false; // couldn't find leaf

                    // update leaf counts
                    for (TrieLeaf tl = (TrieLeaf) n; tl.tuple != t; tl = tl.next)
                        tl.leafCount--;

                    nl.next = nnl.next;
                    b.leafCount--;
                    return true;
                }
            }
        }
    }

    private void repairBranch(TrieBranch b, int i) {
        if (i == 0) {
            b.children[0] = null;
        } else {
            int len = b.chars.length;
            char[] nchars = new char[len - 1];
            TrieNode[] nkids = new TrieNode[len - 1];
            System.arraycopy(b.chars, 0, nchars, 0, i);
            System.arraycopy(b.children, 0, nkids, 0, i);
            System.arraycopy(b.chars, i + 1, nchars, i, len - i - 1);
            System.arraycopy(b.children, i + 1, nkids, i, len - i - 1);
            b.chars = nchars;
            b.children = nkids;
        }
    }

    private void addLeaf(TrieBranch b, TrieLeaf l, int depth) {
        b.leafCount += l.leafCount;

        char c = getChar(l.word, depth);
        int i = getIndex(b.chars, c);

        if (i == -1) {
            addChild(b, l, c);
        } else {
            TrieNode n = b.children[i];
            if (n == null) {
                // we have completely spelled out the word
                b.children[i] = l;
            } else if (n instanceof TrieBranch) {
                // recurse down the tree
                addLeaf((TrieBranch) n, l, depth + 1);
            } else {
                // node is a leaf, need to do a split?
                TrieLeaf nl = (TrieLeaf) n;
                if (i == 0 || (caseSensitive ? nl.word.equals(l.word) : nl.word.equalsIgnoreCase(l.word))) {
                    // same word, so chain the entries
                    for (; nl.next != null; nl = nl.next)
                        nl.leafCount++;
                    nl.leafCount++;
                    nl.next = l;
                } else {
                    // different words, need to do a split
                    TrieBranch nb = new TrieBranch();
                    b.children[i] = nb;
                    addLeaf(nb, nl, depth + 1);
                    addLeaf(nb, l, depth + 1);
                }
            }
        }
    }

    private void addChild(TrieBranch b, TrieNode n, char c) {
        int len = b.chars.length;
        char[] nchars = new char[len + 1];
        TrieNode[] nkids = new TrieNode[len + 1];
        System.arraycopy(b.chars, 0, nchars, 0, len);
        System.arraycopy(b.children, 0, nkids, 0, len);
        nchars[len] = c;
        nkids[len] = n;
        b.chars = nchars;
        b.children = nkids;
    }

    /**
     * Look up the given word in this Trie. If a match is found, a TrieNode
     * is returned. This node is the root of a subtree containing all the
     * matches to the query.
     * @param word the word to lookup
     * @return the TrieNode root of the subtree containing all matches. A
     * null value is returned if no match is found.
     */
    public TrieNode find(String word) {
        return (word.length() < 1 ? null : find(word, root, 0));
    }

    private TrieNode find(String word, TrieBranch b, int depth) {
        char c = getChar(word, depth);
        int i = getIndex(b.chars, c);
        if (i == -1) {
            return null; // not in trie
        } else if (word.length() - 1 == depth) {
            return b.children[i]; // end of search
        } else if (b.children[i] instanceof TrieLeaf) {
            return equalityCheck(word, (TrieLeaf) b.children[i]);
        } else {
            return find(word, (TrieBranch) b.children[i], depth + 1); // recurse
        }
    }

} // end of class Trie
